from abc import ABC, abstractmethod
from typing import Dict, List
import random
import time


# ========== 基础类定义 ==========
class AIAgent(ABC):
    """
    智能体基类（抽象类）
    所有具体角色智能体的基类，定义了智能体的基本结构和行为
    属性:
        name: 智能体名称
        role: 智能体角色
        tools: 智能体可用的工具字典
    """

    def __init__(self, name: str, role: str):
        """
        初始化智能体
        参数:
            name: 智能体名称
            role: 智能体角色
        """
        self.name = name
        self.role = role
        self.tools = {}

    @abstractmethod
    def execute_task(self, task_data: Dict):
        """
        abstractmethod抽象方法：执行具体任务
        子类必须实现此方法以定义具体的任务执行逻辑
        参数:
            task_data: 任务数据字典
        返回:
            任务执行结果
        """
        pass

    def use_tool(self, tool_name: str, params: Dict):
        """
        多态方法：调用工具的统一接口
        参数:
            tool_name: 工具名称
            params: 工具参数字典
        返回:
            工具执行结果
        异常:
            ValueError: 当工具不存在时抛出
        """
        if tool_name in self.tools:
            return self.tools[tool_name](**params)
        raise ValueError(f"Tool {tool_name} not available")


# ========== 工具装饰器 ==========
def register_tool(tool_name: str):
    """
    工具注册装饰器
    这是一个工厂函数，返回一个装饰器，用于注册工具方法并添加日志功能
    参数:
        tool_name: 工具名称，用于日志输出
    返回:
        装饰器函数，用于包装工具方法
    """

    def decorator(func):
        """
        实际的装饰器函数
        参数:
            func: 被装饰的工具方法
        返回:
            包装后的工具方法，添加了日志功能
        """

        def wrapper(self, *args, **kwargs):
            # 打印工具使用日志
            print(f"[{self.name}] 正在使用工具：{tool_name}")
            # 调用原始工具方法并返回结果
            return func(self, *args, **kwargs)

        return wrapper

    return decorator


# ========== 具体角色类 ==========
class Administrator(AIAgent):
    """
    管理员角色（继承自AIAgent）
    负责给智能体分配任务，协调整个工作流程
    工具:
        - 系统监控: 监控系统状态
    """

    def __init__(self, name: str):
        super().__init__(name, "管理员")
        self.tools = {
            "monitor_system": self.monitor_system,
        }
        self.agents = {}

    @register_tool("系统监控")
    def monitor_system(self) -> Dict:
        return {"cpu": random.randint(0, 100), "memory": random.randint(0, 100)}

    def register_agent(self, agent: AIAgent):
        self.agents[agent.role] = agent

    def dispatch_task(self, task_type: str, data: Dict):
        agent = self.agents.get(task_type)
        if agent:
            return agent.execute_task(data)
        raise ValueError(f"No agent for {task_type}")

    def execute_task(self, task_data: Dict):
        print(f"{self.name} 开始执行管理任务")
        if task_data["action"] == "assign_research":
            topic = task_data.get("topic")
            if topic:
                print(
                    {
                        "status": "assigned",
                        "task_type": "research",
                        "topic": topic,
                    }
                )
                return self.dispatch_task("研究员", {"topic": topic})
            else:
                return None
        elif task_data["action"] == "assign_writing":
            outline = task_data.get("outline")
            if outline:
                print(
                    {
                        "status": "assigned",
                        "task_type": "writing",
                        "outline": outline,
                    }
                )
                return self.dispatch_task("撰稿人", {"outline": outline})
            else:
                return None
        elif task_data["action"] == "assign_publishing":
            content = task_data.get("content")
            platform = task_data.get("platform")
            if content and platform:
                print(
                    {
                        "status": "assigned",
                        "task_type": "publishing",
                        "content": content,
                        "platform": platform,
                    }
                )
                return self.dispatch_task(
                    "发布员", {"content": content, "platform": platform}
                )
            else:
                return None
        else:
            return self.use_tool("monitor_system", {})


class Researcher(AIAgent):
    """
    调研员角色（继承自AIAgent）
    负责进行网络搜索和文档解析，收集相关信息
    工具:
        - 网络搜索: 在互联网上搜索指定主题
        - 文档解析: 读取并解析文档内容
    """

    def __init__(self, name: str):
        super().__init__(name, "研究员")
        self.tools = {
            "web_search": self.web_search,
            "read_document": self.read_document,
        }

    @register_tool("网络搜索")
    def web_search(self, query: str) -> List[str]:
        # 模拟网络的搜索延迟
        print(f"[{self.name}] 正在搜索：{query}")
        time.sleep(3)
        # 模拟搜索结果返回
        return [f"{query}的搜索结果{random.randint(1, 100)}" for _ in range(3)]

    @register_tool("文档解析")
    def read_document(self, file_path: str) -> str:
        return f"模拟读取{file_path}内容"

    def execute_task(self, task_data: Dict):
        print(f"{self.name} 开始执行调研任务")
        results = self.use_tool("web_search", {"query": task_data["topic"]})
        return {"status": "success", "data": results}


class Writer(AIAgent):
    """
    撰稿人角色（继承自AIAgent）
    负责根据大纲生成内容并进行格式优化
    工具:
        - 内容生成: 根据大纲生成文章内容
        - 格式优化: 优化文本格式
    """

    def __init__(self, name: str):
        super().__init__(name, "撰稿人")
        self.tools = {
            "generate_content": self.generate_content,
            "format_text": self.format_text,
        }

    @register_tool("内容生成")
    def generate_content(self, outline: str) -> str:
        # 模拟生成内容的延迟
        print(f"[{self.name}] 正在生成内容...")
        time.sleep(5)
        # 模拟生成内容，请替换为实际的大模型生成逻辑
        sections = outline["sections"]
        content = []
        for idx, section in enumerate(sections):
            content.append(f"第{idx + 1}节：{section}")
            content.append(f"内容：{section}")
        return "\n".join(content)

    @register_tool("格式优化")
    def format_text(self, text: str) -> str:
        return text.upper()

    def execute_task(self, task_data: Dict):
        print(f"{self.name} 开始执行写作任务")
        draft = self.use_tool("generate_content", {"outline": task_data["outline"]})
        return {"status": "success", "content": draft}


class Publisher(AIAgent):
    """
    发布员角色（继承自AIAgent）
    负责将内容发布到指定平台或安排定时发布
    工具:
        - 内容发布: 将内容发布到指定平台
        - 定时发布: 安排内容在指定时间发布
    """

    def __init__(self, name: str):
        super().__init__(name, "发布员")
        self.tools = {
            "publish_content": self.publish_content,
            "schedule_post": self.schedule_post,
        }

    @register_tool("内容发布")
    def publish_content(self, content: str, platform: str) -> str:
        # 模拟发布内容的延迟
        print(f"[{self.name}] 正在发布内容到{platform}...")
        time.sleep(2)
        return f"已在{platform}发布内容：{content[:50]}..."

    @register_tool("定时发布")
    def schedule_post(self, content: str, time: str) -> str:
        return f"已安排于{time}发布内容：{content[:50]}..."

    def execute_task(self, task_data: Dict):
        print(f"{self.name} 开始执行发布任务")
        if "schedule_time" in task_data:
            return self.use_tool(
                "schedule_post",
                {"content": task_data["content"], "time": task_data["schedule_time"]},
            )
        else:
            return self.use_tool(
                "publish_content",
                {"content": task_data["content"], "platform": task_data["platform"]},
            )


# ========== 主程序流程 ==========
if __name__ == "__main__":
    # 创建并注册智能体
    admin = Administrator("AdminBot")
    researcher = Researcher("ResearchBot")
    writer = Writer("WriterBot")
    publisher = Publisher("PublishBot")

    admin.register_agent(admin)
    admin.register_agent(researcher)
    admin.register_agent(writer)
    admin.register_agent(publisher)

    # 协作任务流程
    print("AdminBot开始协调任务流程...")

    # 1. 管理员拆解任务并分配给研究员
    print("\n=== 研究阶段 ===")
    research_task = {"action": "assign_research", "topic": "AI发展趋势"}
    research_result = admin.execute_task(research_task)
    print(f"研究结果：{research_result}")

    # 2. 管理员将研究结果交给撰稿员
    print("\n=== 撰稿阶段 ===")
    outline = {"sections": research_result["data"]}
    write_task = {"action": "assign_writing", "outline": outline}
    draft = admin.execute_task(write_task)
    print(f"草稿内容：{draft}")

    # 3. 管理员将稿件交给发布员
    print("\n=== 发布阶段 ===")
    publish_task = {
        "action": "assign_publishing",
        "content": draft["content"],
        "platform": "官网",
    }
    published = admin.execute_task(publish_task)
    print(f"发布结果：{published}")

    print("\n=== 最终输出 ===")
    print(f"[研究结果] {research_result['data']}")
    print(f"[草稿内容] {draft['content']}")
    print(f"[发布结果] {published}")

    print("\n任务流程结束！")
